/*
 * Copyright 2009 Sid Stamm <sid@sidstamm.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * Author:
 * Sid Stamm <sid@sidstamm.com>
 */
Components.utils.import("resource://gre/modules/XPCOMUtils.jsm");

const Cc = Components.classes;
const Ci = Components.interfaces;
const DEBUG_SERVICE = true;

/***********************************************************
class definition
***********************************************************/
function CertLock() {
  this.dump("XPCOM CERTLOCK service creation.");

  this.db = new CertLockDB();

  // set up observer for HTTP requests
  try {
    this.obsService = Cc["@mozilla.org/observer-service;1"]
                        .getService(Ci.nsIObserverService);
    this.obsService.addObserver(this, "http-on-modify-request", false);
    this.obsService.addObserver(this, "http-on-examine-response", false);
    this.obsService.addObserver(this, "http-on-examine-cached-response", false);
  } catch(anException) {
    this.dump("ERROR: couldn't add observer service.");
  }
}


// class definition
CertLock.prototype = {

  DATABASE_NAME:    "certlock_db.sqlite",

  // properties required for XPCOM registration:
  classDescription: "CertLock XPCOM Component",
  classID:          Components.ID("{C39CF9B9-18D3-4CF6-A5C4-5B2CB56E8D9A}"),
  
  contractID:       "@sidstamm.com/certlock;1",
  _xpcom_categories: [{ category: "app-startup", service: true }],
  QueryInterface: XPCOMUtils.generateQI([Ci.ssICertLock, Ci.nsIObserver,
                                         Ci.nsIURIContentListener,
                                         Ci.nsISupportsWeakReference,
                                         Ci.nsISupports]),

  //some debugging functions
  dump: function my_dump(x)
  {
    if(DEBUG_SERVICE) { dump(" CERTLOCK XPC *** " + x + "\n"); }
  },

  dump_props: function my_dump_props(x, showValues) 
  {
    for(let p in x) {
      if(showValues) {
        this.dump(" >> [" + p + "] = " + x[p]);
      } else {
        this.dump(" >> [" + p + "] ");
      }
    }
  },

  init: function(aWindow) {
    //Could add some window-tracking code here, but whatever.
  },

  /**
   * URI content listener.
   */
  onStartURIOpen: function(aURI) {
    this.dump("Opening URI: " + aURI.asciiSpec);
  },
  

  /**
   * observer service callback that watches for HTTPS requests and sends the
   * cert info back to the browser watcher.
   */
  observe: function(subject, topic, data) {
    switch(topic) {
      case "http-on-modify-request":
        var channel = subject.QueryInterface(Ci.nsIHttpChannel);
        //this.dump(">> Request for URI: " + channel.URI.asciiSpec);
        if (!channel.URI.schemeIs("https")) {
          //this.dump("Not HTTPS... punting. ");
          break;
        }
        if (!channel.securityInfo || !(channel.securityInfo instanceof Ci.nsISSLStatusProvider)) {
          //this.dump("No SecurityInfo for channel... punting. " + channel.securityInfo);
          break;
        }
        var cert = channel.securityInfo.QueryInterface(Ci.nsISSLStatusProvider)
                          .SSLStatus.QueryInterface(Ci.nsISSLStatus).serverCert;
        this.dump("CERTIFICATE ORIGIN: " + cert.organization);
        break;

      case "http-on-examine-cached-response":
        this.dump("Examining Cached Response");

      case "http-on-examine-response":
        var response = subject.QueryInterface(Ci.nsIHttpChannel);
        //this.dump(">> Response for URI: " + response.URI.asciiSpec);
        
        // TODO: verify that this works for:
        // 1.  chrome-based XHR?
        // 2.  JS-caused requests (content-based XHR)
        // 3.  JS-caused requests (dom manipulation)
        
        // Process the request.
        if(response.URI.scheme == "https") {
          if(response.securityInfo instanceof Ci.nsISSLStatusProvider) {

            var ssp = response.securityInfo.QueryInterface(Ci.nsISSLStatusProvider);
            var cert = ssp.SSLStatus.QueryInterface(Ci.nsISSLStatus).serverCert;
            var isDocument = !!(response.loadFlags & Ci.nsIHttpChannel.LOAD_DOCUMENT_URI);

            this.dump("HTTPS response:");
            this.dump("        URI: " + response.URI.asciiSpec);
            this.dump("  DOCUMENT?: " + isDocument);
            this.dump("       DATA: " + data);

            // if the channel is not locked, we don't have to enforce
            // certLock.
            if (cert && this.db.channelIsLocked(response.URI.host)) {
            //if (cert && true) {

              // ensure that either the cert is familiar or the CA of the new
              // cert at least comes from the same locked origin (country).
              if (!this.db.certIsFamiliar(response.URI.host, cert)) {
              //if (true) {
                // if it's not familiar, we must cancel the connection.

                // TODO: Should this error be a document?  Do we want to suppress
                // loading of the whole page if an embedded resource's channel violates
                // this problem?

                if (isDocument) {
                  // ONLY redirect requests that are base documents to an error page,
                  // everything else must just be cancelled (no good UI for that
                  // yet).
                  var nCLE = new CertLockEntry(response.URI.host, cert);
                  this.redirectToErrorPage(subject, "oldCertOrigin", nCLE.ca_origin);
                } else
                  response.cancel(Components.results.NS_BINDING_ABORTED);

                return;
              }
            }
          }
        }
        break;
    }
  },

  redirectToErrorPage: function(aRequest, oldOrigin, newOrigin)
  {
    var ou = aRequest.URI.clone();
    var ios = Cc["@mozilla.org/network/io-service;1"].getService(Ci.nsIIOService);
    var ns = "about:lockedcert?oo="+oldOrigin+"&no="+newOrigin+"&u="+aRequest.URI.asciiSpec;
    var nu = ios.newURI(ns, null, null);
    var ir = (aRequest.notificationCallbacks ?
              aRequest.notificationCallbacks :
              aRequest.loadGroup.notificationCallbacks)
             .QueryInterface(Ci.nsIInterfaceRequestor);

    // cancel request and start new request (for error page)
    aRequest.cancel(Components.results.NS_BINDING_ABORTED);
    ir.getInterface(Ci.nsIWebNavigation).loadURI(ns, Ci.nsIWebNavigation.LOAD_FLAGS_BYPASS_HISTORY, null, null, null);
  },

  sslChannelOpen: function(aRequest)
  {
  },

};

/***********************************************************
  The database wrapped structure that takes care of what to
  do when certificates are sighted.
 ***********************************************************/
var CertLockDB = function() {
  //NOTE: TOO EARLY to init db here :( wait to do it later.
  this._cache = {};
};

CertLockDB.prototype = {
  /**
   * Accessor for the database, takes care of initializing it the first time
   * it's accessed.
   */
  get db () {
    if (this._mDBConn != null)
      return this._mDBConn;
    else {
      this._initDB();
      return this._mDBConn;
    }
  },

  /** 
   * This connects to the db and loads all the existing records.  This is
   * called automatically the first time "this.db" is used.
   */
  _initDB:
  function() {
    try {
      // init DB and read stored data
      this._storageService = Cc["@mozilla.org/storage/service;1"]  
                              .getService(Ci.mozIStorageService);  
      var dbfile = Cc["@mozilla.org/file/directory_service;1"]  
                      .getService(Ci.nsIProperties)  
                      .get("ProfD", Ci.nsIFile);  

      CertLock.prototype.dump(" Profile DIR: " + dbfile.path);
      dbfile.append(CertLock.prototype.DATABASE_NAME);  

      // openDatabase also create the file if it does not exist  
      this._mDBConn = this._storageService.openDatabase(dbfile);

      CertLock.prototype.dump(" Inited DB: " + dbfile.path);

      this._mDBConn.executeSimpleSQL("CREATE TABLE IF NOT EXISTS certlock "
                                     + "(host string, "
                                     + "cert_id string, "
                                     + "ca_id string, "
                                     + "ca_origin string) ");
    } catch(e) {
      CertLock.prototype.dump("Error connecting to certlock database: " + e);
      CertLock.prototype.dump(e.stack);
      return;
    }

    //load certlock info from the database
    var entriesToTrash = [];
    var statement = this.db.createStatement("SELECT * FROM certlock");  
    while (statement.executeStep()) {  
      CertLock.prototype.dump("  -> loading record");
      let host = statement.row.host;  
      let cert_id = statement.row.cert_id;  
      let ca_id = statement.row.ca_id;  
      let ca_origin = statement.row.ca_origin;  

      CertLock.prototype.dump(" +++1 " + host + " : " + cert_id + " of " + ca_id + " from " + ca_origin);

      // de-serialize into write-through cache
      let entry = new CertLockEntry(host, cert_id, ca_id, ca_origin);
      this._cache[entry.host] = entry;
    }  
    statement.finalize();
  },

  /**
   * INTERNAL:
   * Adds a database record.
   */
  _asyncAddRecord:
  function _cl_asyncAddRecord(aCertLockEntry, errorCallback) {
    //stores a new record in the database 
    try {
      var statement = this.db.createStatement(
            "INSERT INTO certlock (host, cert_id, ca_id, ca_origin) "
          + "values (:host, :cert_id, :ca_id, :ca_orig)");  
      statement.params.host = aCertLockEntry.host;
      statement.params.cert_id = aCertLockEntry.cert_id;
      statement.params.ca_id = aCertLockEntry.ca_id;
      statement.params.ca_orig = aCertLockEntry.ca_origin;
    } catch(e) {
      CertLock.prototype.dump("ERROR - " + e);
    }

    statement.executeAsync({  
      handleResult: function(aResultSet) {
        CertLock.prototype.dump(" handleResult in _asyncAddRecord()");
      },
      
      handleError: function(aError) {  
        errorCallback("DB ADD Error: " + aError.message);
      },  
      
      handleCompletion: function(aReason) {  
        CertLock.prototype.dump(" handleCompletion in _asyncAddRecord()");
        if (aReason != Ci.mozIStorageStatementCallback.REASON_FINISHED)  
          errorCallback(" DB ADD ERROR: Query canceled or aborted!");  
      }  
    });  
    statement.finalize(); 
  },

  /**
   * INTERNAL:
   * Deletes an existing database record.
   */
  _asyncDelRecord:
  function _ftls_asyncDelRecord(aCertLockEntry, errorCallback) {
    //deletes a record in the database 
    var statement = this.db.createStatement("DELETE FROM certlock WHERE host=:host");  
    statement.params.host = aCertLockEntry.host;

    statement.executeAsync({  
      handleResult: function(aResultSet) {  },
      handleError: function(aError) {  
        errorCallback("DB DELETE Error: " + aError.message);
      },  
      handleCompletion: function(aReason) {  
        if (aReason != Ci.mozIStorageStatementCallback.REASON_FINISHED)  
          errorCallback(" DB DELETE ERROR: Query canceled or aborted!");  
      }  
    });  
    statement.finalize(); 
  },

  /**
   * INTERNAL:
   * Modifies an existing database record.
   */
  _asyncModifyRecord:
  function _ftls_asyncModifyRecord(aCertLockEntry, errorCallback) {
    //stores a new record in the database 
    var upd_sql = "UPDATE certlock "
                + "SET cert_id=:cert_id, "
                + "ca_id=:ca_id, "
                + "ca_origin=:ca_origin "
                + "WHERE host=:host";
    var statement = this.db.createStatement(upd_sql);
    statement.params.host = aCertLockEntry.host;
    statement.params.cert_id = aCertLockEntry.cert_id;
    statement.params.ca_id = aCertLockEntry.ca_id;
    statement.params.ca_origin = aCertLockEntry.ca_origin;

    statement.executeAsync({  
      handleResult: function(aResultSet) {},
      handleError: function(aError) {  
        errorCallback("DB UPDATE Error: " + aError.message);
      },  
      handleCompletion: function(aReason) {  
        if (aReason != Ci.mozIStorageStatementCallback.REASON_FINISHED)  
          errorCallback(" DB UPDATE ERROR: Query canceled or aborted!");  
      }  
    });  
    statement.finalize(); 
  },



  /**
   * Given a CertLockEntry (hostname+cert data), do what is necessary to
   * update it in or add it to the database.
   * IMPORTANT: this will override anything in the database.  If this isn't
   * absolutely necessary, call "certIsFamiliar" first to see if the current
   * channel's cert is ok to use.
   */
  lockCert: 
  function(aCertLockEntry, errorCallback) {
    // Force the db connection to initialize
    this.db;

    var oldEntry = this._cache[aCertLockEntry.host];

    //update the record in the list
    this._cache[aCertLockEntry.host] = aCertLockEntry;

    // .. and the db
    if(oldEntry) {
      CertLock.prototype.dump(" ... modified existing entry");
      return this._asyncModifyRecord(aCertLockEntry, errorCallback);
    }
    CertLock.prototype.dump(" ...  creating new entry");
    return this._asyncAddRecord(aCertLockEntry, errorCallback);
  },

  /**
   * Given a host name and x509 certificate, determines if the new cert is ok.
   * NOTE: this returns false if the site isn't currently certLocked.
   */
  certIsFamiliar:
  function(aHostName, aX509Cert) {
    // Force the db connection to initialize if it hasn't already been loaded.
    this.db;

    if (!this._cache[aHostName.toLowerCase()])
      return false;

    let existing  = this._cache[aHostName.toLowerCase()];
    let potential = new CertLockEntry(aHostName, aX509Cert);

    return existing.lockAllows(potential);
  },

  /**
   * Given a host name, determines if its channels are certLocked.
   */
  channelIsLocked:
  function(aHostName) {
    // force DB connection if it's not already open.
    this.db;
    return !!(this._cache[aHostName.toLowerCase()]);
  }
};

function CertLockEntry(aHostName, aX509Cert) {
  let issuer = aX509Cert.issuer;
  this.host = aHostName.toLowerCase();
  this.cert_id = aX509Cert.md5Fingerprint + aX509Cert.serialNumber;
  this.ca_id = issuer.md5Fingerprint + issuer.serialNumber;
  this.ca_origin = 'unknown';

  let sn = issuer.subjectName.toLowerCase();
  let bits = sn.split(',');
  for (var i in bits) {
    if (bits[i].substring(0,2) == "c=")
      this.ca_origin = bits[i].substring(2);
  }
}

CertLockEntry.prototype = {
  equals: function(that)
  {
    return this.host == that.host
      && this.cert_id == that.cert_id
      && this.ca_id == that.ca_id
      && this.ca_origin == that.ca_origin;
  },

  lockAllows: function(that) 
  {
    return this.ca_origin == that.ca_origin;
  },
};



var components = [CertLock];
function NSGetModule(cmpMgr, fileSpec) {
  return XPCOMUtils.generateModule(components);
}
